#!/usr/local/bin/python
#-*- encoding:utf-8 -*-

import numpy as np;
import scipy.sparse as sp
from multiprocessing import Pool
from itertools import izip_longest
import os
import utils
import sys
import argparse

def loadMean(trainsetflag):
    global D_Mean, Q_Mean;
    D_Mean_Name, Q_Mean_Name = utils.spellMeanSet(trainsetflag);
    D_Mean = np.matrix(np.load(D_Mean_Name));
    print D_Mean.max();
    Q_Mean = np.matrix(np.load(Q_Mean_Name));
    print Q_Mean.max();

def getColumnVector(s, d):
    row = [];
    col = [];
    data = [];
    li = s.strip().split(' ');
    for l in li:
        pair = l.strip().split(':');
        if len(pair) != 2:
            continue;
        index = int(pair[0]);
        value = float(pair[1]);
        row.append(index-1);
        col.append(0);
        data.append(value)
        #data.append(1);
    return sp.csc_matrix((data,(row, col)), shape=(d,1));

    
def score(x, y):
    s1 = x.transpose().dot(M);
    s = s1.dot(y);
    return s;

def score3(x, y):
    x1 = Lx.transpose().dot(x);
    x2 = Ly.transpose().dot(y);

    s1 = x1.transpose().dot(x2).todense();
    s2 = x.transpose().dot(y).todense();
    #print 's1, s2', s1, s2;
    return ((s1 + s2)[0][0], s2[0][0]);
    
def score2(x, y):
    x1 = Lx.transpose().dot(x).todense();
    x2 = Ly.transpose().dot(y).todense();
    
    s1 = x1.transpose().dot(x2);
    s2 = Lxm.transpose().dot(x2)
    s3 = x1.transpose().dot(Lym);
    s5 = (x-Q_Mean.transpose()).transpose().dot(y-D_Mean.transpose());
    s6 = x.transpose().dot(y);
    return ((s1-s2-s3+s4+s5)[0][0], s5[0][0]);

def grouper(n, iterable, padvalue=None):
    return izip_longest(*[iter(iterable)]*n, fillvalue=padvalue);

def processChunk(line):
    if not line:
        return None;
    #print 'pid:', os.getpid();
    li = line.split(',');
    if len(li) != 6:
        print 'Wrong';
    queryID = li[0];
    appID = li[1];
    queryFeature = getColumnVector(li[2], dx);
    documentFeature = getColumnVector(li[3], dy);
    labeled = int(float(li[4]));
    freq = int(int(li[5]));
    #print queryFeature.shape, documentFeature.shape
    queryFeature  = queryFeature;
    documentFeature = documentFeature;
    sm, sb = score3(queryFeature, documentFeature);

    return ('%s\t%f\t%d\t%s\t%d\n' % (queryID, sm, labeled,  appID, freq), '%s\t%f\t%d\t%s\t%d\n' % (queryID, sb, labeled,  appID, freq));
    
def testing(Lx1, Ly1, testset, resultFile, bm25_resultFile):
    global Lx, Ly;
    Lx = Lx1;
    Ly = Ly1;
    test = open(testset, 'r');
    results_file = open(resultFile, 'w');
    bm25_resultFile = open(bm25_resultFile, 'w');
    line = test.readline();
    if not line.startswith('*'):
        print 'No Dimension Provided';
        return;
    global dx, dy;
    firstLi = line.split('\t');
    dx = int(firstLi[1]);
    dy = int(firstLi[2]);
    
    print Lx.shape, Ly.shape
    global Lxm, Lym, s4
    #Lxm = Lx.transpose().dot(Q_Mean.transpose());
    #Lym = Ly.transpose().dot(D_Mean.transpose());
    #print Lxm, Lym;
    #s4 = Lxm.transpose().dot(Lym);
    count = 0;
    K = 20;
    pool = Pool(processes=K);

    for chunk in grouper(K, test):
        results = pool.map(processChunk, chunk);
        for result in results:
                if result:
                    results_file.write(result[0]);
                    bm25_resultFile.write(result[1]);


    results_file.close();
    bm25_resultFile.close();
    test.close();
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Run to get Results from given modelflag trained by trainsetflag on testsetflag');
    parser.add_argument('--modelflag', help= 'the runflag of this run');
    parser.add_argument('--testsetflag', help='specify the testset');
    parser.add_argument('--trainsetflag', help='specify trainset');
    args = parser.parse_args();

    modelflag = args.modelflag;
    trainsetflag = args.trainsetflag;
    testsetflag = args.testsetflag;
    loadMean(trainsetflag);
    LxFile, LyFile = utils.spellModelName(modelflag, trainsetflag);
    testset = utils.spellTestSet(testsetflag);
    resultFile, bm25_result_file = utils.spellResultsName(modelflag, trainsetflag, testsetflag);
    print LxFile, LyFile, resultFile;
    Lx = utils.load_sparse_csr(LxFile);
    Ly = utils.load_sparse_csr(LyFile);
    testing(Lx, Ly, testset, resultFile, bm25_result_file);
